package gateway.config;

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.user.OidcUser;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;

import java.util.Objects;

/**
 * https://docs.spring.io/spring-cloud-gateway/docs/current/reference/html/#naming-custom-filters-and-references-in-configuration
 */
@Component
class IdTokenRelayGatewayFilterFactory extends AbstractGatewayFilterFactory<Object> {
    @Override
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> exchange.getPrincipal()
                .filter(principal -> principal instanceof OAuth2AuthenticationToken)
                .cast(OAuth2AuthenticationToken.class)
                .map(OAuth2AuthenticationToken::getPrincipal)
                .filter(oAuth2User -> oAuth2User instanceof OidcUser)
                .cast(OidcUser.class)
                .map(OidcUser::getIdToken)
                .filter(Objects::nonNull)
                .map(OidcIdToken::getTokenValue)
                .filter(Objects::nonNull)
                .map(idToken -> withBearerAuth(exchange, idToken))
                .defaultIfEmpty(exchange)
                .flatMap(chain::filter);
    }

    ServerWebExchange withBearerAuth(ServerWebExchange exchange, String idToken) {
        return exchange.mutate().request(r -> r.headers(headers -> headers.setBearerAuth(idToken))).build();
    }
}